function [optima, RegionSet, SampleSet, RegionSampleId, OptSetRemain] = PRS_MMO_test_v2( Problem, AlgorithmM, OptSet )
% PRS_MMO: Partition-based Random Search for Multimodal Optimization for 
%   minimization problems.
% decision variable: d dimensions; feature of regions: m dimensions.
%INPUT
%   Problem.
%       Dimension: scalar, the dimension of the decision variables
%       Domain: 1 -by- m vector, features of the domain
%       Partition: function handle of partitioning strategy
%           [RegionNum, Region, RegionSampleId] = Partition(Region, SampleSet)
%               RegionNum: scalar, number of new partitioned regions
%               Region: RegionNum -by- (2+m) matrix,
%                   [partition depth, partitionable, features ...]
%                   partition depth = 0 for the original domain
%                   partitionable: scalar, 1: partitionable, 0: non-partitionable
%               RegionSampleId: RegionNum -by- 1 cell, the sample id that
%                   belongs to each region
%               SampleSet: [SampleID, objective value, x1, x2, ...]
%       Sampling: function handle of sampling in a particular region
%           [Sample] = Sampling(n, region)
%               n: number of samples
%               region: 1 -by- m vector, [features...]
%               Sample: n -by- (1+d) matrix, [objective value, x1, x2, ...]
%   AlgorithmM.
%       n0: scalar, the base sample size
%       QuantileLevel: scalar, <0.5, the quantile level of promisinhg region definition
%       NewBudget: scalar, added budget size at each stage
%       MaxSampleSize: scalar, maximal sample size in one region
%       Radius: scalar, the radius distance that be regarded as neighbor
%       StopCriteria: [X,Y] (the maximal budget size is 1e6)
%           [1,Y]: reach maximum budget Y
%           [2,Y]: reach the required number of optima Y
%           [3,Y]: the number of optima is not changed in Y iterations
%   OptSet.
%       sol: ~ -by- d matrix, the optimal solution
%       fval: scalar, the global optimum
%       epsilon: scalar, the difference of the value considered to be optimum
%       radius: scalar, the distance of the solution considered to be found
%
%OUTPUT
%   optima: ~ -by- (2+d) matrix, the final optimal solution set,
%         [sample id, objective value, x1, x2, ...]
%   RegionSet: ~ -by- (2+m) matrix, the whole subregion set
%         [partition depth, partitionable, features ...]
%   SampleSet: ~ -by- (1+d) matrix, the original sample set
%         [objective value, x1, x2, ...]
%   RegionSampleId: ~ -by- 1 cell, the id of the samples in each region

%% INITIALIZATION
% problem setting
d = Problem.Dimension;
domain = [0, 1, Problem.Domain]; % [partition depth, partitionable, features ...]
Partition = Problem.Partition;
Sampling = Problem.Sampling;
% algorithm setting
n0 = AlgorithmM.n0;
alpha = AlgorithmM.QuantileLevel;
Delta = AlgorithmM.NewBudget;
MaxSampleSize = AlgorithmM.MaxSampleSize;
StopCriteria = AlgorithmM.StopCriteria;
N = (StopCriteria(1)==1)*StopCriteria(2) + (StopCriteria(1)~=1)*1e6; % the maximal budget size
radius = AlgorithmM.radius;
% the optimal set for testing purpose
if ~isempty(OptSet)
    OptSetRemain = OptSet.sol;
end
% algorithm archive
RegionSet = zeros(round(N/n0), length(domain));  % [partition depth, partitionable, features ...]
RegionSet(1,:) = domain;
RegionsNum = 1;  % the number of exist regions
PartitionableNum = 1;  % the number of exist partitionable regions
CurrentMaxPartitionDepth = 0;  % the maximum depth of partitioning at current iteration
RegionSampleId = cell(round(N/n0), 1);  % id of group samples
SampleSet = zeros(N, 2+d);  % [SampleID, objective value, x1, x2, ...]
SampleSet(:,1) = 1:size(SampleSet,1);
N0 = 0;  % budget allocated
OptCandidate = zeros(N,1);  % 1: the possible optima, 0: not possible
GroupInf = zeros(round(N/n0), 6); % [size, mean, std, modified size, quantile, weight]
%% OPTIMIZATION
NumOpt = 0;  % the current optimal value
NumIterationRepeat = 0;  % the number of iterations that the number of optima is not changed
PartitionList = zeros(Delta,1);  % the list of regions to be partitioined
PartitionList(1) = 1;
PartitionListNum = 1;  % the number of the partition list
WeightChangedList = zeros(round(N/n0),1);  % the list of regions whose weight is changed
WeightChangedListNum = 0;  % the number of the region list whose weight is changed
all_weight_changed = 0;  % if all weights should be recalculated
NewOptCandidate = 0;  % if there are new regions reach the non-partitionable size
BestGroupInf = [1,inf,0];  % the best group information
%       [group id, quantile, modified sample size]
% subsequent iterations
while N0 < N && PartitionableNum>0 ...
    && (StopCriteria(1)~=2 || NumOpt<StopCriteria(2)) ...
    && (StopCriteria(1)~=3 || NumOpt==0 || NumIterationRepeat<StopCriteria(2)) ...
    && (isempty(OptSet) || ~isempty(OptSetRemain))
%% partition partitionable regions
    i = 1;
    while i <= PartitionListNum
        % partition
        iPartitionRegion = PartitionList(i);
        [NewRegionNum, NewRegion, NewRegionSampleId]...
            = Partition(RegionSet(iPartitionRegion,:), SampleSet(RegionSampleId{iPartitionRegion},:));
        NewRegionId = [iPartitionRegion,(RegionsNum+1):(RegionsNum+NewRegionNum-1)]; % id of new-partitioned regions
        RegionSet(NewRegionId,:) = NewRegion;
        RegionSampleId(NewRegionId) = NewRegionSampleId;
        RegionsNum = RegionsNum+NewRegionNum-1;  % update total region number
        PartitionableNum = PartitionableNum-1+sum(NewRegion(:,2));  % update partitionable region number
        % update the maximal partition depth
        MaxPartitionDepth_temp = max(NewRegion(:,1));
        if MaxPartitionDepth_temp>CurrentMaxPartitionDepth
            CurrentMaxPartitionDepth = MaxPartitionDepth_temp;
            GroupInf(1:RegionsNum,4) = max(2,round(GroupInf(1:RegionsNum,1).*RegionSet(1:RegionsNum,1)./CurrentMaxPartitionDepth));
            all_weight_changed = 1;
            BestGroupInf(3) = GroupInf(BestGroupInf(1),4);
        end
        % first stage sampling and group information updating
        for j = NewRegionId
            if RegionSet(j,2)==0
                %% non-partitionable region
                if isempty(RegionSampleId{j})
                    N0 = N0+1;
                    SampleSet(N0,2:end) = Sampling(1,RegionSet(j,3:end));
                    RegionSampleId{j} = N0;
                end
                GroupInf(j,1) = length(RegionSampleId{j});
                group_update(j)
                % remove the found optimum
                if ~isempty(OptSet)
                    remove_opt(SampleSet(RegionSampleId{j},2:end));
                end
                OptCandidate(RegionSampleId{j})=1;  % update the candidate optimal solution
                NewOptCandidate = 1;
            else
                %% partitionable region
                GroupInf(j,1) = length(RegionSampleId{j});  % update the sample size
                if GroupInf(j,1)>=MaxSampleSize
                    % if the sample size in the new region also reaches the threshold
                    PartitionListNum = PartitionListNum+1;
                    PartitionList(PartitionListNum) = j;
                else
                    if GroupInf(j,1)<n0
                    % if the sample size in the new region is not enough
                        Add_n = min(n0-GroupInf(j,1),N-N0);
                        NewSampleId = ((N0+1):(N0+Add_n))';
                        SampleSet(NewSampleId,2:end) = Sampling(Add_n,RegionSet(j,3:end));
                        N0 = N0+Add_n;
                        RegionSampleId{j} = [RegionSampleId{j};NewSampleId];
                        GroupInf(j,1) = GroupInf(j,1)+Add_n;
                    end
                    group_update(j)
                    WeightChangedListNum = WeightChangedListNum+1;
                    WeightChangedList(WeightChangedListNum) = j;
                end
            end
        end
        i = i+1;
    end
    PartitionListNum = 0;
%% update the optima set
    if StopCriteria(1)~=1 && NewOptCandidate == 1
        optima = extract();
        NumOpt_old = NumOpt;
        NumOpt = size(optima,1);
        if NumOpt_old ~= NumOpt
            NumIterationRepeat = 0;
        end
        NewOptCandidate = 0;
    end
    NumIterationRepeat = NumIterationRepeat+1;
%% update the weights using baqm
    if all_weight_changed == 1
    % if the maximal partition depth or the minimal quantile is changed
        WeightChangedListNum = PartitionableNum;
        WeightChangedList(1:PartitionableNum) = find(RegionSet(1:RegionsNum,2)==1)';
        all_weight_changed = 0;
    end
    baqm(WeightChangedList(1:WeightChangedListNum))
    WeightChangedListNum = 0;
%% additional budget
    Add_n = min(N-N0,Delta);
    [ n2 ] = add_budget(Add_n);
    for i = 1:size(n2,1)
        iregion = n2(i,1);
        NewSampleId = ((N0+1):(N0+n2(i,2)))';
        SampleSet(NewSampleId,2:end) = Sampling(n2(i,2),RegionSet(iregion,3:end));
        N0 = N0+n2(i,2);
        RegionSampleId{iregion} = [RegionSampleId{iregion};NewSampleId];
        GroupInf(iregion,1) = GroupInf(iregion,1)+n2(i,2);
        if GroupInf(iregion,1)>=MaxSampleSize && RegionSet(iregion,2)==1
            PartitionListNum = PartitionListNum+1;
            PartitionList(PartitionListNum) = iregion;
        else
            group_update(iregion)
            WeightChangedListNum = WeightChangedListNum+1;
            WeightChangedList(WeightChangedListNum) = iregion;
        end
    end
end
%% final results
if StopCriteria(1)==1
%     [optima] = extract();
    optima = [];
end
RegionSet((RegionsNum+1):end, :) = [];
SampleSet((N0+1):end, :) = [];
SampleSet(:,1)=[];
RegionSampleId((RegionsNum+1):end, :) = [];

%% Plots for 2 dimension box constraint problems
% if d==2
%     figure()
%     hold on
%     for i=1:RegionsNum
%         rectangle('Position',[RegionSet(i,[3,5]),RegionSet(i,[4,6])-RegionSet(i,[3,5])])
%     end
%     scatter(SampleSet(:,2),SampleSet(:,3),2,SampleSet(:,1),'filled');
%     hold off
% end


%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% FUNCTIONS
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function baqm(list)
    %baqm: calculate the weight for regions in the list using Budget
    %Allocation for Quantile Minimization
        if ~isempty(list)
            Wk_best = (1+norminv(alpha,0,1)^2-1/BestGroupInf(3));
            K = length(list); % number of groups
            cv_hat = GroupInf(list,3)./(GroupInf(list,2)-BestGroupInf(2));
            Wk = (1+1./(cv_hat.^2)-1./GroupInf(list,4));
            Ckb = Wk_best./Wk;
            P_kb = fcdf(Ckb, GroupInf(list,4)-1, ones(K,1).*BestGroupInf(3)-1); % the probability that cv_k>cv_b
            weights = P_kb./(1-P_kb); % Ni / Nb
            GroupInf(list,6) = weights;
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ n2 ] = add_budget(NewBudget)
    % add_budget: allocate NewBudget according to the weights
    % n2: K -by- 2 matrix, [region id, new sample size]
        n1 = GroupInf(1:RegionsNum,4);
        N_tot = sum(n1)+NewBudget;
        n = N_tot.*GroupInf(1:RegionsNum,6)./sum(GroupInf(1:RegionsNum,6));
        n2 = zeros(RegionsNum,1);
        while NewBudget>0
            [n_delta, k] = max(n-n1-n2);
            if n_delta<1
                add = 1;
            else
                add = max(0, min(NewBudget,round(n_delta)));
            end
            n2(k) = n2(k)+add;
            NewBudget = NewBudget-add;
        end
        n2_id = find(n2>0);
        n2 = [n2_id, n2(n2_id)];
%         n2 = [(1:RegionsNum)',n2];
%         n2(n2(:,2)<=0,:)=[];
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function [ Optima ] = extract()
    %Nitching: obtain the final optima set
        %% ranking the solutions
        [~, Ranking] = sort(SampleSet(1:N0,2),'descend');
        SampleSet_temp = SampleSet(Ranking,:); % sort the samples from bad to good
        OptimaSet = OptCandidate(Ranking,:);
        %% clearing
        sampleid = find(OptimaSet,1);
        while 1
            Distance = (sqrt(sum((repmat(SampleSet_temp(sampleid,3:end),[N0,1])-SampleSet_temp(:,3:end)).^2,2)))';
            neighbor = find(Distance<=radius);
            [~,best_neighbor] = min(SampleSet_temp(neighbor,2));
            OptimaSet(neighbor) = 0; % delete all neighbor
            OptimaSet(neighbor(best_neighbor)) = 1; % save the best neighbor to the optima set
            next = find(OptimaSet((sampleid+1):end),1);
            if isempty(next)
                break
            else
                sampleid = sampleid + next;
            end
        end
        Optima = SampleSet_temp(OptimaSet==1,:);
        OptCandidate = zeros(N,1);
        OptCandidate(Optima(:,1))=1;
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function remove_opt( new_samples )
    % remove_opt: remove the found optimal solution from OptSetRemain
        new_samples(new_samples(:,1)>=(OptSet.fval+OptSet.epsilon),:) = [];
        for ii = 1:size(new_samples,1)
            OptDis = sqrt(sum((repmat(new_samples(ii,2:end),[size(OptSetRemain,1),1]) - OptSetRemain).^2,2));
            [MinDis, OptId] = min(OptDis);
            if MinDis<OptSet.radius
                OptSetRemain(OptId,:) = [];
            end
        end
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    function group_update( groupId )
    % update the group information
        Samplevalue_temp = SampleSet(RegionSampleId{groupId},2);
        % update group information
        GroupInf(groupId,2) = mean(Samplevalue_temp);  % group mean
        GroupInf(groupId,3) = sqrt(var(Samplevalue_temp));  % group variance
        GroupInf(groupId,4) = max([2,round(GroupInf(groupId,1)*RegionSet(groupId,1)/CurrentMaxPartitionDepth)]);  % modified sample size
        new_quantile = GroupInf(groupId,2)+GroupInf(groupId,3)*norminv(alpha,0,1);
        GroupInf(groupId,5) = new_quantile;
        if groupId==BestGroupInf(1) || new_quantile<BestGroupInf(2)
            if new_quantile<=BestGroupInf(2)
                BestGroupInf = [groupId, new_quantile, GroupInf(groupId,4)];
            else
                [BestGroupInf(2),BestGroupInf(1)] = min(GroupInf(1:RegionsNum,5));
                BestGroupInf(3) = GroupInf(BestGroupInf(1),4);
            end
            all_weight_changed = 1;
        end
        GroupInf(groupId,6) = 0;  % the weight
    end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
end

